package com.zk.algorithm.tree;

import com.zk.algorithm.annotation.Easy;
import com.zk.algorithm.annotation.LeetCodeExplore;
import com.zk.algorithm.bean.TreeNode;

import java.util.*;

/**
 * @author zk
 */
@LeetCodeExplore
@Easy
public class LevelOrderTraversal {

    // ====================
    // 当前位于第几层，使用 Map 存储
    // ====================

    public List<List<Integer>> levelOrder(TreeNode root) {
        List<List<Integer>> mList = new ArrayList<>();
        if (root == null) {
            return mList;
        }

        Queue<TreeNode> mQueue = new LinkedList<>();
        Map<TreeNode, Integer> mLevel = new HashMap<>();

        mQueue.offer(root);
        mLevel.put(root, 0);

        int cur = -1;
        while (!mQueue.isEmpty()) {
            TreeNode mNode = mQueue.poll();
            int level = mLevel.get(mNode);

            if (level != cur) {
                List<Integer> row = new ArrayList<>();
                row.add(mNode.val);
                mList.add(row);
                cur = level;
            } else {
                mList.get(level).add(mNode.val);
            }

            if (mNode.left != null) {
                mQueue.offer(mNode.left);
                mLevel.put(mNode.left, cur + 1);
            }

            if (mNode.right != null) {
                mQueue.offer(mNode.right);
                mLevel.put(mNode.right, cur + 1);
            }
        }

        return mList;
    }

    // ====================
    // 但其实当前位于第几层，完全可以使用 queue.size() 来确定
    //           3
    //          / \
    //         9  20
    //           /  \
    //          15   7
    // ====================

    public List<List<Integer>> levelOrder2(TreeNode root) {
        Queue<TreeNode> queue = new LinkedList<TreeNode>();
        List<List<Integer>> wrapList = new LinkedList<List<Integer>>();

        if (root == null) return wrapList;

        queue.offer(root);
        while (!queue.isEmpty()) {
            // [3] => 1
            // [9, 20] => 2
            // [15, 17] => 2
            int levelNum = queue.size();
            List<Integer> subList = new LinkedList<Integer>();

            for (int i = 0; i < levelNum; i++) {
                TreeNode cur = queue.poll();
                subList.add(cur.val);

                if (cur.left != null)
                    queue.offer(cur.left);

                if (cur.right != null)
                    queue.offer(cur.right);
            }

            wrapList.add(subList);
        }

        return wrapList;
    }

}
